"""
@project:StrategyBacktestEngine
@author:liuzeyu
@time:2022-09-08
@version:v0.0.1
"""
import tushare as ts
import efinance
import pandas as pd
from datetime import datetime


# 使用tushare获取数据，根据backtrader的数据格式对数据进行处理
def get_data_tushare(code, start='2010-01-01', end='2020-07-16', export: bool = True):
    df = ts.get_k_data(code, autype='qfq', start=start, end=end)
    df.index = pd.to_datetime(df.date)
    df['openinterest'] = 0
    if export:
        pd.DataFrame(df).to_csv(path_or_buf=str(code + '-' + start + 'to' + end + '.csv'), index=False)
    return df


# 使用efinance获取数据，根据backtrader的数据格式对数据进行处理
def get_data_efinance(stock_code, begin: datetime, end: datetime) -> pd.DataFrame:
    """
    根据efinance工具包获取股票数据
    :param stock_code:股票代码
    :param begin: 开始日期
    :param end: 结束日期
    :return:
    """
    # stock_code = '600519'  # 股票代码，茅台
    k_dataframe: pd.DataFrame = efinance.stock.get_quote_history(
        stock_code, beg=begin.strftime("%Y%m%d"), end=end.strftime("%Y%m%d"))
    k_dataframe = k_dataframe.iloc[:, :9]
    k_dataframe.columns = ['name', 'code', 'date', 'open', 'close', 'high', 'low', 'volume', 'turnover']
    k_dataframe.index = pd.to_datetime(k_dataframe.date)
    k_dataframe.drop(['name', 'code', "date"], axis=1, inplace=True)
    return k_dataframe



